package com.zhxq.ai.websocket;

import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.config.annotation.EnableWebSocket;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;

@Slf4j
@Component
@EnableWebSocket
@ServerEndpoint("/member/ws/tts/{chatId}")
public class TTSServer {

    private static Map<String, Session> sessionMap = new HashMap<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("chatId") String chatId) {
        log.info("WebSocketCallServer onOpen" + session.getId());
        sessionMap.put(chatId, session);
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("接收到了客户端 {} 发来的消息 : {}", session.getId(), message);
    }

    /**
     * 流式输入输出音频
     * @param chatId
     * @param audioFrame
     */
    public void sendAudio(String chatId, ByteBuffer audioFrame) {
        byte[] audioBytes = new byte[audioFrame.remaining()];
        audioFrame.get(audioBytes);
        Session session = sessionMap.get(chatId);
        if (session != null && session.isOpen()) {
            try {
                // 使用 java.nio.ByteBuffer.wrap 包装 bytes 成 Binary Frame
                session.getBasicRemote().sendBinary(ByteBuffer.wrap(audioBytes));
            } catch (IOException e) {
                log.error("发送合并后的音频失败", e);
            }
        }
    }

    /**
     * 发送结束标识
     * @param chatId
     */
    public void sendEndSignal(String chatId) {
        Session session = sessionMap.get(chatId);
        if (session != null && session.isOpen()) {
            try {
                // 发送文本 "EOS" 表示音频流结束
                session.getBasicRemote().sendText("EOS");
            } catch (IOException e) {
                log.error("发送结束标识失败", e);
            }
        }
    }

    public void closeSession(String chatId) {
        Session session = sessionMap.get(chatId);
        if (session != null && session.isOpen()) {
            try {
                session.close();
            } catch (IOException e) {
                log.warn("关闭 session 失败", e);
            }
        }
        sessionMap.remove(chatId);
    }
}
